source("HNC_metadata_tidy.R")
source("risk_factor_regressions.R")

SBSburden <- read.csv("Input_mutational_burden/Input_SBS_mutational_burden.csv") %>% rename(SBS_burden = Mutational.burden)
DBSburden <- read.csv("Input_mutational_burden/Input_DBS_mutational_burden.csv") %>% rename(DBS_burden = Mutational.burden)
IDburden <- read.csv("Input_mutational_burden/Input_ID_mutational_burden.csv") %>% rename(ID_burden = Mutational.burden)

ancestry <- read.csv("Supplementary_Table_5_data_ancestry.csv") %>% select(-country)

additional_df <- list(SBSburden, DBSburden, IDburden, ancestry)

for (n in 1:length(additional_df)) {
  data <- data %>% left_join(additional_df[[n]], by = "donor_id")
}

data <- data %>%
  mutate(
    country = factor(country, levels = c("Brazil", "Colombia", "Argentina", "Greece", "Italy", "Czech Republic", "Slovakia", "Romania")),
    age = scale(age_diag),
    subsite = factor(subsite, levels = c("Larynx", "OC", "OPC", "Hypopharynx"))
  )

signatures <- grep("_burden", names(data), value = T)
factors <- c(
  "sex", "stage", "age",
  "OC", "OPC", "Larynx", "Hypopharynx",
  "region", names(ancestry)[-c(1, 2)],
  "alcohol_ever", "tobacco_ever",
  "hpv_pos_opc", "bmi_cat", "oral_cat", "hotdrinks"
)
confounders <- c("age", "sex", "region", "subsite", "alcohol_ever", "tobacco_ever")

SBS_limit <- 100000
DBS_limit <- 6000
ID_limit <- 5000

output <- NULL
for (f in factors) {
  for (s in signatures) {
    # remove hypermutated
    if (s == "SBS_burden") {
      subset <- data %>% filter(SBS_burden < SBS_limit)
    } else if (s == "DBS_burden") {
      subset <- data %>% filter(DBS_burden < DBS_limit)
    } else if (s == "ID_burden") {
      subset <- data %>% filter(ID_burden < ID_limit)
    }

    # remove cases with missing data
    subset <- subset[!is.na(subset[, f]) & !is.na(subset[, s]), ]

    # remove confounders that overlap with the parameter of interest
    confounders_minus <- confounders
    if (f %in% c(
      "country", "incidence", "argentina", "greece", "italy", "brazil", "czechRepublic", "romania", "predicted_ancestry",
      "European", "African", "East_Asian", "Amerindian"
    )) {
      confounders_minus <- confounders_minus[confounders_minus != "region"]
    }
    if (f %in% c("hpv_pos_opc", "OC", "OPC", "Larynx", "Hypopharynx")) {
      confounders_minus <- confounders_minus[!confounders_minus %in% c("subsite")]
    }
    if (f %in% c("tobacco", "heavy_smoker", "heavy_smoker_dic", "years_since_stop_tob_cat")) {
      confounders_minus <- confounders_minus[confounders_minus != "tobacco_ever"]
    }
    if (f %in% c("alcohol", "heavy_drinker", "heavy_drinker_dic", "years_since_stop_alc_cat")) {
      confounders_minus <- confounders_minus[confounders_minus != "alcohol_ever"]
    }
    if (f %in% c("tob_alc", "known_riskfactor")) {
      confounders_minus <- confounders_minus[!confounders_minus %in% c("tobacco_ever", "alcohol_ever")]
    }
    # remove counfonders with only 1 level in the analyzed df
    for (c in confounders_minus[confounders_minus != "age"]) { ## age is a continuous variable so should not be taken into account for this
      if (ncol(table(as.character(subset[, f]), as.character(subset[, c]))) < 2) {
        print(paste("** Warning: Confounding", c, "has only one level for parameter", p, "in signature", s, "and will be removed"))
        confounders_minus <- confounders_minus[confounders_minus != c]
      }
    }

    # Linear regression
    myformula <- paste0(s, " ~ ", paste(c(f, confounders_minus), collapse = " + "))
    print(myformula)
    model <- lm(myformula, data = subset)
    result <- summary(model)$coefficients
    r <- as.data.frame(result) %>%
      tibble::rownames_to_column("independent_vars") %>%
      select(-`t value`) %>%
      filter(grepl(f, independent_vars)) %>%
      rename(pval = `Pr(>|t|)`) %>%
      mutate(p_adj = ifelse(pval * length(signatures) < 1, pval * length(signatures), 1), .after = "pval") %>%
      mutate(mutation_burden = s, .after = "independent_vars")
    output <- rbind(output, r)
  }
}

write.csv(output, "output/Supplementary_Table_5.csv", row.names = F)
